{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hll/anaconda3/envs/tensorflow/lib/python3.6/site-packages/sklearn/preprocessing/_encoders.py:415: FutureWarning: The handling of integer data will change in version 0.22. Currently, the categories are determined based on the range [0, max(values)], while in the future they will be determined based on the unique values.\n",
      "If you want the future behaviour and silence this warning, you can specify \"categories='auto'\".\n",
      "In case you used a LabelEncoder before this OneHotEncoder to convert the categories to integers, then you can now use the OneHotEncoder directly.\n",
      "  warnings.warn(msg, FutureWarning)\n",
      "/home/hll/anaconda3/envs/tensorflow/lib/python3.6/site-packages/sklearn/preprocessing/_encoders.py:415: FutureWarning: The handling of integer data will change in version 0.22. Currently, the categories are determined based on the range [0, max(values)], while in the future they will be determined based on the unique values.\n",
      "If you want the future behaviour and silence this warning, you can specify \"categories='auto'\".\n",
      "In case you used a LabelEncoder before this OneHotEncoder to convert the categories to integers, then you can now use the OneHotEncoder directly.\n",
      "  warnings.warn(msg, FutureWarning)\n",
      "/home/hll/anaconda3/envs/tensorflow/lib/python3.6/site-packages/sklearn/preprocessing/_encoders.py:415: FutureWarning: The handling of integer data will change in version 0.22. Currently, the categories are determined based on the range [0, max(values)], while in the future they will be determined based on the unique values.\n",
      "If you want the future behaviour and silence this warning, you can specify \"categories='auto'\".\n",
      "In case you used a LabelEncoder before this OneHotEncoder to convert the categories to integers, then you can now use the OneHotEncoder directly.\n",
      "  warnings.warn(msg, FutureWarning)\n",
      "/home/hll/anaconda3/envs/tensorflow/lib/python3.6/site-packages/sklearn/preprocessing/_encoders.py:415: FutureWarning: The handling of integer data will change in version 0.22. Currently, the categories are determined based on the range [0, max(values)], while in the future they will be determined based on the unique values.\n",
      "If you want the future behaviour and silence this warning, you can specify \"categories='auto'\".\n",
      "In case you used a LabelEncoder before this OneHotEncoder to convert the categories to integers, then you can now use the OneHotEncoder directly.\n",
      "  warnings.warn(msg, FutureWarning)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train: (560000, 204)\n",
      "train: (560000, 204)\n",
      "train: (560000, 204)\n",
      "train: (352035, 204)\n",
      "test: (140001, 204)\n",
      "test: (140001, 204)\n",
      "test: (140001, 204)\n",
      "test: (88009, 204)\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import LabelEncoder, OneHotEncoder\n",
    "from sklearn.preprocessing import MinMaxScaler, StandardScaler\n",
    "\n",
    "#Mark the nominal columns and consolidate the data (extract all the nominal columns)\n",
    "#标出名词性列，整合数据（把名词性的列都提取出来）\n",
    "def combine_dataset(files, col_names, processed = False):  \n",
    "\tdtypes = {}\n",
    "\tif processed == False:\n",
    "\t\tfor col_name in col_names:\n",
    "\t\t\tnominal_names = set(['srcip', 'sport', 'dstip', 'dsport', 'proto', 'state',\n",
    "\t\t\t\t                 'service', 'ct_ftp', 'label_10'])  #Nominal column\n",
    "\t\t\tif col_name in nominal_names:\n",
    "\t\t\t\tdtypes[col_name] =  str\n",
    "\t\t\telse:\n",
    "\t\t\t\tdtypes[col_name] = np.float32\n",
    "\telse:\n",
    "\t\tfor col_name in col_names:\n",
    "\t\t\tdtypes[col_name] = np.float32\n",
    "\n",
    "\trecords = []\n",
    "\tfor file in files:\n",
    "\t\tdata = pd.read_csv(file, header = None, names = col_names, dtype = dtypes)\n",
    "\t\trecords.append(data)\n",
    "\n",
    "\trecords_all = pd.concat(records)#当没有索引时、concat不管列名，直接加到一起\n",
    "                                     #When there is no index, concat adds them together regardless of the column names,\n",
    "                     \n",
    "\n",
    "\treturn records_all\n",
    "\n",
    "## Make new col names for categorical features after one-hot encoding\n",
    "#为one-hot编码之后的列起个新列名\n",
    "def get_nominal_names(dataset, cols_nominal):  \n",
    "\tdata_nominal = dataset[cols_nominal]  \n",
    "\n",
    "\n",
    "\n",
    "\tnew_col_names = []\n",
    "\tfor col_name in cols_nominal:\n",
    "\t\tname_unique = sorted(dataset[col_name].unique())  #名词性列的不同的值。Different values for noun columns\n",
    "\t\tnew_col_name = [col_name + '_' + x for x in name_unique]  \n",
    "\t\tnew_col_names.extend(new_col_name)\n",
    "\n",
    "\treturn new_col_names\n",
    "\n",
    "#Remove the unimportant feature, one-hot encoding, and convert the attack class to numeric\n",
    "#删除不重要的特征，one-hot编码，将攻击类别转换为数值型\n",
    "def select_feature_and_encoding(dataset, cols_to_drop, cols_nominal, cols_nominal_all):\n",
    "    \n",
    "\n",
    "\t# Drop the features has no meaning such as src ip. 删除不重要的特征\n",
    "\tfor cols in cols_to_drop:\n",
    "\t\tdataset.drop(cols, axis = 1, inplace = True)\n",
    "\n",
    "\t# Save the label and then drop it from dataset 保留标签然后将它从数据集中删除（提取出标签列）\n",
    "\tlabel_10 = dataset['label_10']\n",
    "\tlabel_2 = dataset['label_2']\n",
    "\tdataset.drop('label_2', axis = 1, inplace = True)\n",
    "\tdataset.drop('label_10', axis = 1, inplace = True)\n",
    "\n",
    "\t# replace the label with specific code  将标签数值化\n",
    "\treplace_dict = { np.nan: 0, 'Analysis': 1, 'Backdoors': 2, 'Backdoor': 2, 'DoS': 3,\n",
    "                    'Exploits':4,' Fuzzers': 5, ' Fuzzers ':5, 'Generic': 6,\n",
    "                    'Reconnaissance': 7, ' Shellcode ':8, 'Shellcode': 8,\n",
    "                    'Worms':9, ' Reconnaissance ': 7,}\n",
    "\tnew_label_10 = label_10.replace(replace_dict)\n",
    "\tnew_label_10.to_frame() \n",
    "\tlabel_2.to_frame\n",
    "\tdel label_10\n",
    "\n",
    "\t# replace the lost values  用0替换缺失值\n",
    "\treplace_dict = {np.nan: 0, ' ': 0}\n",
    "\tfor cols in ['ct_ftp', 'ct_flw', 'is_ftp']:\n",
    "\t\tdataset[cols] = dataset[cols].replace(replace_dict)\n",
    "\n",
    "\t# 'is_ftp' column is wrong, correct it(I found that the value of it is\n",
    "\t# all the same with ct_ftp_cmd, so if the value is not 0, is_ftp should\n",
    "\t# be 1)\n",
    "\tfor x in dataset['is_ftp']:\n",
    "\t\tif x != 0:\n",
    "\t\t\tx = 1\n",
    "\n",
    "\t# select and process the categorical features 选择并处理分类特征\n",
    "\tdata_nominal = dataset[cols_nominal]  #cols_nominal是名词性列的列名，提取出名词性列的数据\n",
    "\tdata_temp_1 = data_nominal.apply(LabelEncoder().fit_transform)  #将名词性列进行编号\n",
    "\tdel data_nominal\n",
    "\n",
    "\n",
    "\tnew_col_names = []\n",
    "\tfor col_name in cols_nominal:\n",
    "\t\tname_unique = sorted(dataset[col_name].unique())\n",
    "\t\tnew_col_name = [col_name + '_' + x for x in name_unique]\n",
    "\n",
    "\t\tnew_col_names.extend(new_col_name)\n",
    "\t\tdataset.drop(col_name, axis = 1, inplace = True) \n",
    "\n",
    "\t#one-hot\n",
    "\tenc = OneHotEncoder()\n",
    "\tdata_temp_2 = enc.fit_transform(data_temp_1)\n",
    "\tdel data_temp_1 \n",
    "\n",
    "\tdata_encoded = pd.DataFrame(data_temp_2.toarray(), columns = new_col_names)\n",
    "\tdel data_temp_2\n",
    "\n",
    "\t# complement the nominal columns 补充名词性列\n",
    "\tdiff = set(cols_nominal_all) - set(new_col_names) \n",
    "\n",
    "\tif diff:\n",
    "\t\tfor cols in diff:\n",
    "\t\t\tdata_encoded[cols] = 0.\n",
    "\t\tdata_encoded = data_encoded[cols_nominal_all]\n",
    "\n",
    "\tdataset= dataset.join(data_encoded)  \n",
    "\tdel data_encoded\n",
    "\n",
    "\tdataset = dataset.join(new_label_10)\n",
    "\tdataset = dataset.join(label_2)\n",
    "\n",
    "\treturn dataset  #Complete data set (including data and labels)\n",
    "                    #完整的数据集（包括数据和标签）\n",
    "\n",
    "#Split the training set and test set and save the file as a CSV file\n",
    "#分裂训练集和测试集,并将文件保存成CSV文件\n",
    "def split_dataset(dataset, file_train, file_test):   \n",
    "\n",
    "\tcols = dataset.columns\n",
    "\t#trainset, testset = train_test_split(dataset, test_size = 0.2)\n",
    "\ttrainset, testset = train_test_split(dataset, test_size = 0.2,random_state=40,stratify=dataset['label_10'])\n",
    "\ttrain = pd.DataFrame(trainset, columns = cols)\n",
    "\ttest = pd.DataFrame(testset, columns = cols)\n",
    "\n",
    "\ttrain.to_csv(file_train)\n",
    "\ttest.to_csv(file_test)\n",
    "\n",
    "#Standardize, and save the file in CSV and tf formats\n",
    "#标准化，并将文件保存成csv格式和tf格式\n",
    "def scaling(files_train, files_test, col_names_scaling, scaling_type):\n",
    "\n",
    "\tif scaling_type == 'min_max':\n",
    "\t\tscaler = MinMaxScaler()\n",
    "\t\tfile_folder = 'min_max/'\n",
    "\telse:\n",
    "\t\tscaler = StandardScaler()\n",
    "\t\tfile_folder = 'normalized/'\n",
    "\n",
    "\tif not os.path.exists(file_folder):\n",
    "\t\tos.mkdir(file_folder)\n",
    "\tcols = []\n",
    "\tfor file in files_train:\n",
    "\t\t# col 0 is the index in the file\n",
    "\t\ttrainset = pd.read_csv(file, index_col = 0, dtype = np.float32)\n",
    "\t\tif len(cols) == 0:\n",
    "\t\t\tcols = trainset.columns\n",
    "\t\tscaler.partial_fit(trainset[col_names_scaling])\n",
    "\n",
    "\tdel trainset\n",
    "\tcols_keep = list(set(cols) - set(col_names_scaling))\n",
    "\n",
    "\tfor file in files_train:\n",
    "\t\ttrainset = pd.read_csv(file, dtype = np.float32)\n",
    "\t\ttrain_scaled = scaler.transform(trainset[col_names_scaling])\n",
    "\t\ttrain_changed = pd.DataFrame(train_scaled, columns = col_names_scaling)\n",
    "\t\ttrain_unchanged = trainset[cols_keep]\n",
    "\t\ttrainset_final = pd.concat((train_changed, train_unchanged),\n",
    "\t\t                        axis = 1)\n",
    "\t\ttrainset_final = trainset_final[cols]\n",
    "\t\tprint(\"train:\",trainset_final.shape)  #trainset shape\n",
    "\t\tfile_csv = file_folder + file\n",
    "\t\ttrainset.to_csv(file_csv, index = False)\n",
    "\t\tlen_tail = len('.csv')   \n",
    "\t\tfile_tfr = file_folder + file[:-1 * len_tail] + '.tfrecords'\n",
    "\t\tmake_tfrecords(trainset_final, file_tfr)\n",
    "\n",
    "\tfor file in files_test:\n",
    "\t\ttestset = pd.read_csv(file, dtype = np.float32)\n",
    "\t\ttest_scaled = scaler.transform(testset[col_names_scaling])\n",
    "\t\ttest_changed = pd.DataFrame(test_scaled, columns = col_names_scaling)\n",
    "\t\ttest_unchanged = testset[cols_keep]\n",
    "\t\ttestset_final = pd.concat((test_changed, test_unchanged),axis = 1)\n",
    "\t\ttestset_final = testset_final[cols]\n",
    "\t\tprint(\"test:\",testset_final.shape)\n",
    "\t\tfile_csv = file_folder + file\n",
    "\t\ttestset.to_csv(file_csv, index = False)\n",
    "\t\tlen_tail = len('.csv')\n",
    "\t\tfile_tfr = file_folder + file[:-1 * len_tail] + '.tfrecords'\n",
    "\t\tmake_tfrecords(testset_final, file_tfr)\n",
    "\n",
    "#Save the file in tf format\n",
    "#将文件保存成tf格式\n",
    "def make_tfrecords(dataset, file_to_save):  \n",
    "\n",
    "\ttry:\n",
    "\t\tdata = dataset.values\n",
    "\texcept:\n",
    "\t\tdata = dataset\n",
    "\twith tf.python_io.TFRecordWriter(file_to_save) as writer:\n",
    "\t\tfor rows in data:\n",
    "\t\t\tfeatures, label_10, label_2 = rows[:-2], rows[-2], rows[-1]\n",
    "\t\t\tfeature = {'features': tf.train.Feature(float_list = tf.train.FloatList(value = features)),\n",
    "\t\t\t           'label_2': tf.train.Feature(float_list = tf.train.FloatList(value = [label_2])),\n",
    "\t\t\t           'label_10': tf.train.Feature(float_list = tf.train.FloatList(value = [label_10]))}\n",
    "\t\t\texample = tf.train.Example(features = tf.train.Features(feature = feature))\n",
    "\t\t\twriter.write(example.SerializeToString())\n",
    "\n",
    "def next_batch(filename, batch_size):\n",
    "\n",
    "\tlen_feature = 202  #特征数（不包含标签）。 Number of features (not including tags)\n",
    "\tlen_label = 1#标签长度。 The length of the label\n",
    "\n",
    "\tdef read_data(examples):\n",
    "\t\tfeatures = {\"features\": tf.FixedLenFeature([len_feature], tf.float32),\n",
    "                    \"label_2\": tf.FixedLenFeature([len_label], tf.float32),\n",
    "                    \"label_10\": tf.FixedLenFeature([len_label], tf.float32)}\n",
    "\t\tparsed_features = tf.parse_single_example(examples, features)\n",
    "\t\treturn parsed_features['features'], parsed_features['label_2'], \\\n",
    "               parsed_features['label_10']\n",
    "\n",
    "\tdata = tf.data.TFRecordDataset(filename)\n",
    "\tdata = data.map(read_data)\n",
    "\tdata = data.batch(batch_size)\n",
    "\titerator = data.make_one_shot_iterator()\n",
    "\tnext_data, next_label_2, next_label_10 = iterator.get_next()\n",
    "\n",
    "\treturn next_data, next_label_10, next_label_2\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "\n",
    "\tfile_folder = '/home/hll/IDS/UNSW-NB15 - CSV Files/'  #读取的原始文件所在的位置。 The location where the original file was read\n",
    "\tcol_names = ['srcip', 'sport', 'dstip', 'dsport', 'proto', 'state', 'dur',\n",
    "\t             'sbytes', 'dbytes', 'sttl', 'dttl', 'sloss', 'dloss',\n",
    "\t             'service', 'sload', 'dload', 'spkts', 'dpkts', 'swin', 'dwin',\n",
    "\t             'stcpb', 'dtcpb', 'smeansz', 'dmeansz', 'trans_depth',\n",
    "\t             'res_bdy_len', 'sjit', 'djit', 'stime', 'ltime', 'sintpkt',\n",
    "\t             'dintpkt', 'tcprtt', 'synack', 'ackdat', 'is_sm_ips',\n",
    "\t             'ct_state_ttl', 'ct_flw', 'is_ftp', 'ct_ftp', 'ct_srv_src',\n",
    "\t             'ct_srv_dst', 'ct_dst_ltm', 'ct_src_ltm', 'ct_src_dport',\n",
    "\t             'ct_dst_sport', 'ct_dst_src', 'label_10', 'label_2']    #特证名（列名）。 listed name\n",
    "\n",
    "\tcols_to_drop = ['srcip', 'dstip', 'stime', 'ltime', 'sport', 'dsport'] \n",
    "\tcols_nominal = ['proto', 'service', 'state']   #名词性特征。Nominal features\n",
    "\n",
    "\tfiles = [file_folder + 'UNSW-NB15_' + str(i+1) + '.csv' for i in range(4)]  \n",
    "\tdataset = combine_dataset(files, col_names)   \n",
    "\tcols_nominal_all = get_nominal_names(dataset, cols_nominal)  \n",
    "\tdel dataset  \n",
    "\n",
    "\tfile_tail = len('.csv')  \n",
    "\tfile_head = len(file_folder + 'UNSW-NB15_')   \n",
    "\tdtypes = {}   \n",
    "\tfor col_name in col_names:\n",
    "\t\tnominal_names = set(['srcip', 'sport', 'dstip', 'dsport', 'proto', 'state',\n",
    "\t\t\t                 'service', 'is_ftp', 'ct_flw', 'ct_ftp', 'label_10'])\n",
    "\t\tif col_name in nominal_names:\n",
    "\t\t\tdtypes[col_name] =  str  \n",
    "\t\telse:\n",
    "\t\t\tdtypes[col_name] = np.float32 \n",
    "\n",
    "\tfor file in files:\n",
    "\t\tfile_train = file[file_head:-1 * file_tail] + '_train.csv'  #每个文件分裂出的训练集和测试集，csv文件。 \n",
    "        #Each file is split out of the training set and test set, CSV file\n",
    "\t\tfile_test = file[file_head: -1 * file_tail] + '_test.csv'\n",
    "\t\tdataset = pd.read_csv(file, header = None, names = col_names, dtype = dtypes)  \n",
    "\t\tdataset = select_feature_and_encoding(dataset, cols_to_drop, cols_nominal,\n",
    "\t\t                                          cols_nominal_all)  \n",
    "\t\tsplit_dataset(dataset, file_train, file_test)   \n",
    "\n",
    "\tcols_unchanged = ['is_ftp', 'is_sm_ips'] + cols_nominal +\\\n",
    "\t                 cols_to_drop + ['label_2', 'label_10']   \n",
    "\tcols_scaling = [x for x in col_names if x not in cols_unchanged]  \n",
    "\n",
    "\tfiles_train = [str(x + 1) + '_train.csv' for x in range(4)]  \n",
    "\tfiles_test = [str(x + 1) + '_test.csv' for x in range(4)]  \n",
    "\n",
    "\tscaling(files_train, files_test, cols_scaling, 'std')  #标准化。standardized\n",
    "\n",
    "\tfile_folder = 'normalized/' #数据标准化后存放的文件夹。A folder where data is stored after standardization\n",
    "\tfiles_train = [file_folder + str(x + 1) + '_train.tfrecords' for x in range(4)]  \n",
    "\tfiles_test = [file_folder + str(x + 1) + '_test.tfrecords' for x in range(4)]   \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Integrate the four separate data sets\n",
    "#将分开的4个数据集整合到一起\n",
    "def make_whole_datasets(tfrecords_train, num_train_example, tfrecords_test,\n",
    "                        num_test_example):          \n",
    "\n",
    "    data_train, label_10_train, label_2_train= next_batch(tfrecords_train,num_train_example)\n",
    "    data_test, label_10_test, label_2_test= next_batch(tfrecords_test,num_test_example)\n",
    "    with tf.Session() as sess:\n",
    "        data, label_10, label_2 = sess.run([data_train, label_10_train,label_2_train])\n",
    "    dataset = np.concatenate([data, label_10, label_2], axis = 1)\n",
    "\n",
    "    #trainset, valiset = train_test_split(dataset, test_size = 254004,stratify=dataset['label_10'])  \n",
    "    trainset, valiset = train_test_split(dataset, test_size = 0.125,random_state=40,stratify=dataset[:,-2])\n",
    "    print(\"train:\",trainset.shape)\n",
    "    print(\"val:\",valiset.shape) \n",
    "\n",
    "    make_tfrecords(trainset, 'normalized/train.tfrecords') \n",
    "    make_tfrecords(valiset, 'normalized/validation.tfrecords')\n",
    "\n",
    "    del trainset, valiset\n",
    "\n",
    "    with tf.Session() as sess:\n",
    "        data, label_10, label_2 = sess.run([data_test, label_10_test,label_2_test])\n",
    "    dataset = np.concatenate([data, label_10, label_2], axis = 1)\n",
    "    print(\"test:\",dataset.shape)  \n",
    "    make_tfrecords(dataset, 'normalized/test.tfrecords')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train: (1778030, 204)\n",
      "val: (254005, 204)\n",
      "test: (508012, 204)\n"
     ]
    }
   ],
   "source": [
    "num_train_example = 2032035 #trainset size\n",
    "num_test_example = 508012 #testset size \n",
    "make_whole_datasets(files_train, num_train_example, files_test, num_test_example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
